import pymongo
from pymongo import MongoClient
import random

from IPProxyPool.settings import MONGODB_URL
from IPProxyPool.utils.log import logger
from IPProxyPool.domain import Proxy


class MongoProxiesPool(object):

    def __init__(self):
        # 建立连接
        self.client = MongoClient(MONGODB_URL)
        # 获取操作集合
        self.proxies = self.client['proxy_pool']['proxies']

    def __del__(self):
        # 关闭链接
        self.client.close()

    def insert_one(self, proxy):
        # 实现插入功能
        count = self.proxies.count_documents({'_id': proxy.ip})
        if count == 0:
            # 插入
            # 获取字典样式的数据
            dic = proxy.__dict__
            dic['_id'] = proxy.ip
            self.proxies.insert_one(dic)
            logger.info('插入新代理：{}'.format(proxy))
        else:
            logger.warning('已存在的代理：{}'.format(proxy))

    def update_one(self, proxy):
        # 实现修改功能
        # 按照ip找到对应的 改变里面发生改变的一切
        self.proxies.update_one({'_id': proxy.ip}, {'$set': proxy.__dict__})
        logger.info('已更新代理：{}'.format(proxy))

    def delete_one(self, proxy):
        # 实现删除功能
        self.proxies.delete_one({'_id': proxy.ip})
        logger.info('已删除代理ip：{}'.format(proxy))

    def find_all(self):
        cursor = self.proxies.find()
        for item in cursor:
            # Proxy类里面没有这个选项
            item.pop('_id')

            proxy = Proxy(**item)
            yield proxy

    def find(self, conditions={}, count=0):
        '''
        实现查询功能: 根据条件进行查询, 可以指定查询数量, 先分数降序, 速度升序排, 保证优质的代理IP在上面.
        :param conditions: 查询条件
        :param count: 查询几个
        :return: 返回满足条件的代理列表
        '''

        cursor = self.proxies.find(conditions, limit=count).sort([
            ('score', pymongo.DESCENDING), ('speed', pymongo.ASCENDING)
        ])

        proxy_list = []
        for item in cursor:
            item.pop('_id')
            proxy = Proxy(**item)
            proxy_list.append(proxy)

        return proxy_list

    def get_proxies(self, protocol=None, domain=None, count=0, nick_type=0):
        '''
        实现根据协议类型 和 要访问网站的域名, 获取代理IP列表
        :param protocol:协议
        :param domain:域名
        :param count:限制数量，默认全部
        :param nick_type:匿名类型，默认高匿
        :return:满足的代理
        '''

        # 定义查询条件
        conditions = {'nick_type': nick_type}
        #         根据协议定义查询条件
        if protocol is None:
            # 如果没传，就是两个都支持的代理
            conditions['protocol'] = 2
        elif protocol.lower() == 'http':
            conditions['protocol'] = {'$in': [0, 2]}
        else:
            conditions['protocol'] = {'$in': [1, 2]}
        if domain:
            conditions['disable_domains'] = {'$nin': [domain]}

        return self.find(conditions, count)

    def random_proxy(self, protocol=None, domain=None, count=0, nick_type=0):
        '''
        实现根据协议类型 和 要访问完整的域名, 随机获取一个代理IP
        :param protocol:协议
        :param domain:域名
        :param count:限制数量，默认全部,就是多少个代理当中获取一个，越少性能越高，随机性越差
        :param nick_type:匿名类型，默认高匿
        :return:满足的随机的一个代理
        '''

        # 获取满足条件的所有
        proxies_list = self.get_proxies(protocol=protocol, domain=domain, count=count, nick_type=nick_type)
        # 随机娶一个
        return random.choice(proxies_list)

    def disable_domain(self, ip, domain):
        '''
        根据ip把domain插入到禁止列表
        :param ip:id 地址
        :param domain:域名
        :return:true成功，false失败
        '''

        is_exist = self.proxies.count_documents({'_id': ip, 'disable_domains': domain})
        if is_exist == 0:
            self.proxies.update_one({'_id': ip}, {'$push': {'disable_domains': domain}})
            logger.info('ip：{}禁止域名：{}添加成功'.format(ip, domain))
            return True
        else:
            logger.warning('ip：{}禁止域名：{}已经存在'.format(ip, domain))
            return False


if __name__ == '__main__':
    mongo = MongoProxiesPool()
    a = mongo.get_proxies('http')
    print(a)
    del mongo
